// Copyright (C) 2010 and onwards Noam Shazeer and Georges Harik
// Pieces of this code may be licenced under the Apache License, version 2.0 
//   from Harik Shazeer Labs

#ifndef _ALLOCATOR_H_
#define _ALLOCATOR_H_

#include <stdarg.h>
#include <pthread.h>
#include "util.h"
#include "thing.h"
#include "jobqueue.h"
// #include "table-structures.h"
#include "sync.h"
#include "file.h"
#include "nolock.h"

struct GlobalAllocator;
extern GlobalAllocator * G_ALLOC;

struct LocalAllocatorException {};

void InitAllocator();

// Number of words for global boxed objects
extern const uint64 G_BOX_SPACE;

const uint64 KILL_THREADS_URGENCY = 0;
const uint64 KILL_PROCESS_URGENCY = 20;

extern Sync * G_SYNC;

struct Frontier;

class AllocatorDependant {
 public:
  AllocatorDependant() {}
  virtual ~AllocatorDependant() {}
  virtual int32 NumShards() = 0;
  virtual void AboutToCollect() {}
  virtual void FlagReferences(int32 shard, Frontier * frontier) = 0;
  virtual bool AddToGlobalMap() { return true; }
};

struct ThreadAllocator;

// A set of Objects to be expanded
// The set is a union of 2 types of Object sets
// Object set 1 is a linked link using the link_ field
// Object set 2 is a range inside a big array

struct Frontier {
Frontier() : head_((Box *)0x1), list_expansion_size_(0), 
    start_(NULL), end_(NULL),
    seen_memory_(0), seen_boxes_(0) {}
  void Add(Thing o);

  // Add the range of Objects into the linked list
  void ExpandRange(const Thing * begin, const Thing* end);

  void Expand();

  void AddToQueue(Box *b);

  // For the linked list
  Box * head_;

  // Sum of the array sizes in the linked list
  uint64 list_expansion_size_;
  
  // Everything below is used in load balancing, not yet implemented. 

  // object set 2. 
  Thing * start_;
  Thing * end_;

  // How much memory this frontier has seen
  uint64 seen_memory_;
  uint64 seen_boxes_;
};

// The Global allocator has a big block, and allocates subblocks
// for each ThreadAllocator

struct AllocationBlock {
  uint64 capacity_;  // Words available not including overhead
  uint64 used_;     

  AllocationBlock() { CHECK(false); }

  Memblock StoredMemblock() { 
    return Memblock((char *)this, (Overhead() + used_) * sizeof(uint64)); }

  void Clear() { used_ = 0; }

  void SetAllMemory(uint64 w) { 
    for (uint64 i=0; i<capacity_; i++) Data()[i] = w; }

  void SetUsedMemory(uint64 w) { 
    for (uint64 i=0; i<used_; i++) Data()[i] = w; }

  // Make a AllocationBlock in a memory region
  // Pass in the total free space (including the overhead)
  static AllocationBlock * Make(uint64 *start, uint64 block_size) {
    AllocationBlock * ret = (AllocationBlock *) start;
    ret->capacity_ = block_size - Overhead();
    ret->used_ = 0;
    return ret;
  }
  static uint64 Overhead() { 
    return sizeof(AllocationBlock) / sizeof (uint64); }
  uint64 * Data() { return ((uint64 *)this) + Overhead(); }
  uint64 * DataEnd() { return Data() + used_; }
  AllocationBlock * Next() { 
    return (AllocationBlock *)(Data() + capacity_); }
  uint64 *Allocate(uint64 length) {
    if (used_ + length > capacity_) return NULL;
    uint64 * ret = Data() + used_;
    used_ += length;
    return ret;
  }
  AllocationBlock * AllocateSubblock(uint64 subblock_capacity) {
    uint64 * space = Allocate(subblock_capacity + Overhead());
    if (!space) return NULL;
    AllocationBlock * ret = Make(space, subblock_capacity + Overhead());
    return ret;
  }
};

struct GlobalAllocator {
  GlobalAllocator();

  // Allocate a block. Return NULL on failure
  AllocationBlock * AllocateBlock(uint64 length, bool force = false);

  // Heuristic determining if we've freed enough memory
  // after garbage collection
  bool HasEnoughMemory();

  uint64 AccurateRunningMemoryUsed();

  void ReIndex(bool clear = true);
  void ClearBoxIndex(); // zero out the box index.
  void GarbageCollect(uint64 initial_urgency = 0);
  void AnnounceAboutToCollect();
  void ComputeReachable();
  void MoveAllocationBlocks();

  void Store(string fn);
  void Load(string fn);

  void AddDependant(AllocatorDependant *d) { 
    MutexLocker ml(&lock_); dependants_.insert(d); }
  void RemoveDependant(AllocatorDependant *d) { 
    MutexLocker ml(&lock_); dependants_.erase(d); }

  // The allocator's block
  AllocationBlock * block_;

  // To synchronize allocation
  pthread_mutex_t lock_;
  bool needs_garbage_collection_;

  // These are involved in flagging memory usage
  set<AllocatorDependant *> dependants_;

  // In garbage collection, a job for propagating reachability
  JobQueue::Job *propagate_job_;

  // This value increases as we try to garbage collect more stuff
  uint64 urgency_;

  // Tracked during garbage collection
  uint64 memory_used_;
  uint64 boxes_used_;

  // Configuration options
  uint64 box_space_;
  uint64 emergency_box_space_;
  uint64 block_size_;
  };

enum AllocationMode {
  GLOBAL_ALLOCATION,
  LOCAL_ALLOCATION,
};

// One of these exists per thread that runs a virtual machine
struct ThreadAllocator : public Sync::Worker {
  ThreadAllocator();
  ~ThreadAllocator();

  AllocationBlock *global_block_;
  AllocationBlock *local_block_;
  // uint64 generation_;
  SimpleCache cache_;
  Thing temp_[100]; // to help make small arrays

  RNG rng_;
  uint64 startup_block_size_;
  uint64 global_words_allocated_;

  void ReallocateLocalBlock(uint64 size) {
    if (local_block_) {
      CHECK(local_block_->used_ == 0);
      delete[] (uint64*)local_block_;
    }
    uint64 *p = new uint64[size];
    local_block_ = AllocationBlock::Make(p, size);
  }

  void ClearLocalBlock() { 
    // local_block_->SetUsedMemory(generation_);
    // generation_++;
    local_block_->Clear(); 
    cache_.Clear();
  }
  
  // allocate raw memory - length is in words
  uint64 * GlobalAllocate(uint64 length);
  uint64 * LocalAllocate(uint64 length);

  bool LocalMemoryOK(Thing t) {
    if (t.IsNot()) return true;
    if (t.GetBasicType() != DIRECT_ARRAY) return true;
    uint64* allocation_ptr = local_block_->Data() + local_block_->used_;
    return ((uint64*)t.data_ < allocation_ptr) 
      && ((uint64*)t.data_ >= local_block_->Data());
  }

  void IndexBox(Box *b); // index and cause garbage collection if too full.

  // Thing MakeLocalMutableOrImmutableArray(Thing *contents, uint64 length);

  Thing MakeImmutableArray(AllocationMode m, Thing* contents, uint64 length);
  Thing MakeMutableArray(Thing *contents, uint64 length);
  Thing MakeMutableArray();
  Thing MakeSymbol(const string & s);
  Thing Globalize(Thing t);

  // resizing MArrays.  Reallocates if necessary.
  void ResizeMutableArray(Box * a, uint64 new_length);
  Thing MakeMutableArray(uint64 length, Thing initial_element);

  Thing MakeImmutableArray(AllocationMode m, uint64 length, ...) {
    va_list va;
    va_start(va, length);
    for (uint64 i=0; i<length; i++)
      temp_[i] = va_arg(va, uint64);
    va_end(va);
    return MakeImmutableArray(m, temp_, length);
  }
  Thing MakeImmutableArray0(AllocationMode m) {
    return MakeImmutableArray(m, 0); }
  Thing MakeImmutableArray1(AllocationMode m, Thing t1) {
    return MakeImmutableArray(m, 1, t1.data_); }
  Thing MakeImmutableArray2(AllocationMode m, Thing t1, Thing t2) {
    return MakeImmutableArray(m, 2, t1.data_, t2.data_); }
  Thing MakeImmutableArray3(AllocationMode m, Thing t1, Thing t2, Thing t3) {
    return MakeImmutableArray(m, 3, t1.data_, t2.data_, t3.data_); }
  Thing MakeImmutableArray4(AllocationMode m, Thing t1, Thing t2, Thing t3, 
			    Thing t4) {
    return MakeImmutableArray(m, 4, t1.data_, t2.data_, t3.data_, t4.data_); }
  Thing MakeImmutableArray5(AllocationMode m, Thing t1, Thing t2, Thing t3, 
			    Thing t4, Thing t5) {
    return MakeImmutableArray(m, 5, t1.data_, t2.data_, t3.data_, t4.data_, 
			      t5.data_);}  
  
  // Make an immutable array based on a vector of Things
  Thing MakeImmutableArray(AllocationMode m, vector<Thing>& v) {
    if (v.size() == 0) return MakeImmutableArray(m, (Thing *)NULL, 0);
    return MakeImmutableArray(m, &(v[0]), v.size()); }

  // Same thing on a StableVector
  Thing MakeImmutableArray(AllocationMode m, StableVector<Thing>& v) {
    return MakeImmutableArray(m, &(v[0]), v.size()); }
};

extern string ToString(Thing);

// GARBAGE COLLECTION
void * GarbageCollectMaintenance(void * junk);

// Run a handler on all boxes in a block
struct BlockTask : public JobQueue::Task {
  virtual void Handle(Box * o) {};
  virtual void GlobalStart() {}
  virtual void StartBlock() {}
  virtual void FinishBlock() {}

  void Run() {
    StartBlock();
    Box * o = (Box *) block_->Data();
    for (; o < (Box *) (block_->DataEnd());) {
      Box * next = o->Next();
      if (o->GetType() != FREE_SPACE) Handle(o);
      o = next;
    }
    FinishBlock();
  }

  // Computed internally
  AllocationBlock * block_;
};

// BT is a subclass of BlockTask
template <class BT> 
void RipAllBoxes() {
  JobQueue::Job j(G_JQ);
  bool first = true;
  for (AllocationBlock * b = (AllocationBlock *)(G_ALLOC->block_->Data()); 
       b < (AllocationBlock *)(G_ALLOC->block_->DataEnd()); 
       b = b->Next()) {
    BT * bt = new BT();
    if (first) {
      bt->GlobalStart();
      first = false;
    }
    bt->block_ = b;
    j.AddTask(bt);
  }
  CHECK(!first);
  j.Wait();
}

#endif
